Skip to content

[logging] Split rollout trajectory time into inference-engine wait vs env.step#1900

Open
eicherseiji wants to merge 14 commits into
NovaSky-AI:mainfrom
eicherseiji:seiji/rollout-llm-vs-env-time-metrics
Open

[logging] Split rollout trajectory time into inference-engine wait vs env.step#1900
eicherseiji wants to merge 14 commits into
NovaSky-AI:mainfrom
eicherseiji:seiji/rollout-llm-vs-env-time-metrics

Conversation

@eicherseiji

@eicherseiji eicherseiji commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Splits each trajectory's rollout time into engine time and env time. Today only the total exists (trajectory_completion_time_*, #1804), so a slow rollout could be engine-bound or env-bound and the metrics cannot tell which.

New metrics:

  • generate/trajectory_llm_time_{mean,p90,max}: time in the inference-engine call, summed over turns
  • generate/trajectory_env_time_{mean,p90,max}: time in env.step(), summed over turns
  • generate/trajectory_overhead_time_{mean,p90,max}: everything else in e2e_time (env construction and teardown, tokenization, chat templating, event-loop scheduling), so the three bands sum to the total and unattributed time is visible

How

agent_loop accumulates one time_splits dict per trajectory (keys llm, env). It travels as a single trajectory_time_splits dict-of-lists on GeneratorOutput. Stats are recomputed from the raw lists in concatenate_generator_outputs, which every logging path goes through, because per-group aggregates like a p90 cannot be combined. Overhead is e2e minus every recorded split component.

Scope

#1925 (env setup time band) is stacked on this PR.

Tests

  • test_llm_vs_env_time_split_metrics (new): runs generate() with an env deliberately slower than the engine and asserts each sleep lands in the right component.
  • Exact-value tests for the metric math and its concatenation are in test_generator_output_utils.py. Step-wise replication and the GeneratorOutput field guardrail are asserted in existing tests.
  • uv run --isolated --extra dev --extra skyrl-train pytest tests/train/generators/ passes, 51 tests. ruff and black clean on changed files.

eicherseiji and others added 2 commits July 14, 2026 20:38
… env.step

Rollout timing is currently reported only as a total per trajectory
(`generate/trajectory_completion_time_*`, added in NovaSky-AI#1804). That total cannot
distinguish an engine-bound rollout from an environment-bound one, which is the
first question to ask when GPU utilization sags during generation.

Accumulate, per trajectory in `agent_loop`, the wall-clock time spent awaiting
`inference_engine_client.generate()` and the time spent in `env.step()`, and
report:

  generate/trajectory_llm_time_{mean,p90,max}
  generate/trajectory_env_time_{mean,p90,max}
  generate/frac_time_in_env

`frac_time_in_env` is time-weighted (sum over sum) so long trajectories count
proportionally rather than each trajectory contributing equally.

The two components need not sum to `e2e_time`; the remainder is tokenization,
chat-template rendering and other in-loop bookkeeping. Both are plumbed like the
existing `e2e_time`: optional, and omitted entirely if any trajectory in the
batch did not record them.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… gauge, buffer depth

Three observability additions to the fully-async RL trainer, all motivated by a
post-mortem where reconstructing GPU idle time required correlating Prometheus
(GPU/disk, wall-clock keyed) against W&B (phase timings, step keyed) by hand.

1. Wire RayGpuMonitor into the async loop. It is constructed in the base
   RayPPOTrainer and flushed in the base loop, but FullyAsyncRayPPOTrainer
   overrides train() and never started or flushed it -- so runs on the async
   trainer logged zero `ray/` GPU keys despite `enable_ray_gpu_monitor=True`.
   Start it at loop entry, flush per step into the committed payload, stop it in
   the finally.

2. Publish a `skyrl_training_phase` gauge via `ray.util.metrics`
   (skyrl/train/utils/phase_metrics.py). Ray exports these to the same
   Prometheus that scrapes node GPU metrics, so
   `avg(ray_node_gpus_utilization) by (Phase)` becomes a single-store query
   instead of a manual cross-tracker correlation -- and it works after a cluster
   restart, when the experiment tracker may be unreachable. Set at each existing
   Timer() phase boundary (waiting_for_buffer / converting / training /
   weight_sync / eval / checkpoint); best-effort, never raises.

3. Log `async/gen_buffer_qsize` and `async/gen_buffer_maxsize` at step start.
   Run-ahead depth was previously only in a tqdm postfix; as a metric it
   distinguishes a throughput-limited generator (buffer low) from a
   gate-limited one (buffer at the staleness-bounded maxsize).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@CharlieFRuan CharlieFRuan changed the title [logging] Split rollout trajectory time into inference-engine wait vs env.step [logging] Rollout time split + async-trainer GPU/phase/buffer observability Jul 15, 2026
Item 1's llm/env time split was computed only at the per-generate()
get_rollout_metrics call. Every logging path (async trainer, eval,
step-wise) re-aggregates through concatenate_generator_outputs, which had
no raw llm/env lists and fell back to the substring-based extra_keys
aggregator. There p90 and frac_time_in_env hit the sum() branch, so a
16-group mini-batch logged frac_time_in_env near 12 and p90 near 16x its
real value.

Store the raw per-trajectory llm/env lists on GeneratorOutput like
e2e_time, thread them through concatenate_generator_outputs with the same
step-wise last-step filter, and recompute all seven metrics there so the
extra_keys fallback leaves them alone.

Also in this commit:
- Extend test_llm_vs_env_time_split_metrics to assert the split on a
  concatenated output, the path that actually broke.
- Fix the phase_metrics module docstring query. The gauge has a lowercase
  `phase` tag and no shared label with the GPU metric, so
  `avg(ray_node_gpus_utilization) by (Phase)` does not run; replace it
  with a vector match against `ray_skyrl_training_phase{phase="eval"}`.
- Deduplicate the last-step filter and the frac-block array construction.
Add ScalarGauges, a best-effort ray.util.metrics scalar-gauge helper that
no-ops without Ray, like TrainingPhaseGauge. Publish the async run-ahead
buffer levels through it so buffer pressure sits in the same Prometheus
store as the training-phase gauge and node GPU metrics, not only in the
experiment tracker:

  skyrl_gen_buffer_qsize     completed groups buffered at step start
  skyrl_gen_buffer_maxsize   staleness-bounded buffer capacity
  skyrl_mini_batch_size      groups consumed per training step
  skyrl_gen_group_keep_rate  kept / (kept + dropped) for the last
                             mini-batch, the zero-variance drop rate
                             under sample_full_batch
Rename the run-ahead phrasing to plain buffer depth in the buffer-metric and gauge comments; no behavior change.
The async-trainer observability items bundled here (RayGpuMonitor
wiring, training-phase gauge, buffer depth) move to their own PRs so
each is a reviewable unit. This restores fully_async_trainer.py to base
and removes phase_metrics.py and its tests. The rollout LLM vs env time
split and its aggregation across concatenate_generator_outputs stay.
@eicherseiji eicherseiji changed the title [logging] Rollout time split + async-trainer GPU/phase/buffer observability [logging] Split rollout trajectory time into inference-engine wait vs env.step Jul 17, 2026
SumanthRH pushed a commit that referenced this pull request Jul 17, 2026
)

## What does this PR do?

Fixes a wiring bug: `RayGpuMonitor` never ran on the fully async
trainer.

`RayGpuMonitor` (#1712) is constructed in `RayPPOTrainer.__init__` and
enabled by default with `trainer.enable_ray_gpu_monitor=True`. The base
loop starts it, flushes it into the per-step log payload, and stops it.
`FullyAsyncRayPPOTrainer` overrides `train()` and never called any of
the three, so fully async runs logged no `ray/` GPU keys at all.
Confirmed on real runs.

The async loop now starts the monitor at loop entry, flushes it into the
per-step timing payload, and stops it in the `finally`, mirroring the
base loop wiring.

## Tests

Wiring only, covered by lint and existing tests. Exercising it end to
end needs a GPU cluster. A structural lifecycle test that catches any
base-loop hook dropped by an overriding trainer is a planned follow-up.

Split out of #1900.
Review feedback pass:
- Add trajectory_llm_times/env_times to the GeneratorOutput field
  guardrail test, which this branch had broken.
- Extract _add_time_stats for the mean/p90/max blocks (completion, llm,
  env) and reuse its sums for frac_time_in_env.
- Extract _optional_times for the omit-if-any-None per-prompt lists.
- Use _concat_field for stop_reasons and rollout_logprobs too.
- Cut duplicated comments; the llm/env split was explained in six
  places, now documented once per surface.
- Assert the step-wise per-step replication of the split in the
  existing step-wise timing test.
Review feedback: the llm and env times traveled as parallel named fields
(llm_time/env_time on two dataclasses, trajectory_llm_times/
trajectory_env_times on GeneratorOutput), so every hop re-documented and
re-plumbed each field, and the omit-if-any-None handling needed its own
helper per field.

One time_splits dict per trajectory (keys "llm", "env") and one
trajectory_time_splits dict-of-lists on GeneratorOutput replace them.
The split semantics are documented once, on TrajectoryOutput.time_splits.
_optional_times is gone; get_rollout_metrics loops the dict through
_add_time_stats generically, so a future component is a new key rather
than a new field at every layer. Metric names are unchanged.
frac_time_in_env is derivable exactly as env_mean / (llm_mean + env_mean)
from metrics already logged, so it added convenience, not information.

Replace the "need not sum to e2e" caveat with the metric that closes the
stack: generate/trajectory_overhead_time_{mean,p90,max}, computed as e2e
minus every recorded split component. Generic over the split keys, so a
future component reallocates time out of overhead without touching this
code. The test asserts overhead is the exact remainder and survives
concatenation.
The generate() harness test conflated timing attribution with metric
math. Attribution keeps the integration test, now reading the splits
straight off GeneratorOutput instead of spying on get_rollout_metrics.
The metric math and concatenation move to exact-value tests over
hand-built outputs in test_generator_output_utils.py. Also drop the
overhead code comment.
The dataclass fields carry no comment; base.py states the keys and that the whole field is None (never partially populated) when any trajectory did not record its split.
@eicherseiji
eicherseiji marked this pull request as ready for review July 18, 2026 00:26

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces tracking and aggregation of trajectory generation time splits (e.g., distinguishing between LLM and environment step times) to compute detailed completion and overhead metrics. The review feedback highlights three potential runtime issues in the newly added utility functions: a potential KeyError in _split_lists due to mismatched keys, a potential TypeError in _concat_field if subsequent outputs contain None values, and a potential TypeError in _concat_time_splits if any trajectory lacks time splits. Robust code suggestions were provided to safely handle these edge cases.

Comment thread skyrl/train/generators/skyrl_gym_generator.py
Comment thread skyrl/train/generators/utils.py Outdated
Comment thread skyrl/train/generators/utils.py Outdated
_concat_time_splits keyed off the first output, so concatenating a batch
that recorded splits with one that did not crashed when the missing
batch came second. Fold list and dict-of-lists handling into one
_concat_optional_field with any-missing-is-None semantics for every
optional field, and let _last_step_only recurse into dicts. Adds an
order-independence regression test.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant